{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nIndividual Populus Trees detection\\n~~~~~~~~~~~~~~~~\\ncode by wHy\\nGhent University\\nHaoyu.Wang@ugent.be\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"\n",
    "Individual Populus Trees detection\n",
    "~~~~~~~~~~~~~~~~\n",
    "code by wHy\n",
    "Ghent University\n",
    "Haoyu.Wang@ugent.be\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gdal\n",
    "import ogr\n",
    "import os\n",
    "import osr\n",
    "import cv2\n",
    "import math\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import fnmatch\n",
    "from collections import deque\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "from pylab import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_img(out_path, im_proj, im_geotrans, im_data):\n",
    "    \"\"\"output img\n",
    "    \"\"\"\n",
    "    if 'int8' in im_data.dtype.name:\n",
    "        datatype = gdal.GDT_Byte\n",
    "    elif 'int16' in im_data.dtype.name:\n",
    "        datatype = gdal.GDT_UInt16\n",
    "    else:\n",
    "        datatype = gdal.GDT_Float32\n",
    "\n",
    "    if len(im_data.shape) > 2:  \n",
    "        im_bands, im_height, im_width = im_data.shape\n",
    "    else:  \n",
    "        im_bands, (im_height, im_width) = 1, im_data.shape\n",
    "\n",
    "    driver = gdal.GetDriverByName(\"GTiff\")\n",
    "    new_dataset = driver.Create(\n",
    "        out_path, im_width, im_height, im_bands, datatype)\n",
    "    new_dataset.SetGeoTransform(im_geotrans)\n",
    "    new_dataset.SetProjection(im_proj)\n",
    "    if im_bands == 1:\n",
    "        new_dataset.GetRasterBand(1).WriteArray(im_data.squeeze())\n",
    "    else:\n",
    "        for i in range(im_bands):\n",
    "            new_dataset.GetRasterBand(i + 1).WriteArray(im_data[i])\n",
    "\n",
    "    del new_dataset\n",
    "\n",
    "def read_img(sr_img):\n",
    "    \"\"\"read raster image\n",
    "    \"\"\"\n",
    "    im_dataset = gdal.Open(sr_img)\n",
    "    if im_dataset == None:\n",
    "        print('open sr_img false')\n",
    "        sys.exit(1)\n",
    "    im_geotrans = im_dataset.GetGeoTransform()\n",
    "    im_proj = im_dataset.GetProjection()\n",
    "    im_width = im_dataset.RasterXSize\n",
    "    im_height = im_dataset.RasterYSize\n",
    "    im_bands = im_dataset.RasterCount \n",
    "    im_data = im_dataset.ReadAsArray(0, 0, im_width, im_height)\n",
    "    del im_dataset\n",
    "\n",
    "    return im_data, im_proj, im_geotrans, im_height, im_width, im_bands\n",
    "\n",
    "def imagexy2geo(trans, row, col):\n",
    "    \"\"\"xy coord to geo coord\n",
    "    \"\"\"\n",
    "    px = trans[0] + col * trans[1] + row * trans[2]\n",
    "    py = trans[3] + col * trans[4] + row * trans[5]\n",
    "    return px, py\n",
    "\n",
    "def write_point_to_layer(coord, out_lyr, def_out_feature):\n",
    "    \"\"\"write point features in layer\n",
    "    \"\"\"\n",
    "    point = ogr.Geometry(ogr.wkbPoint)\n",
    "    point.AddPoint(coord[0], coord[1])\n",
    "    outfeat = ogr.Feature(def_out_feature)\n",
    "    outfeat.SetGeometry(point)\n",
    "    outfeat.SetField2('type', coord[2])\n",
    "    out_lyr.CreateFeature(outfeat)\n",
    "    outfeat = None\n",
    "\n",
    "def formatted_write_opencv_img(img_data, write_path, img_name, prefix, suffix, im_proj, im_geotrans):\n",
    "    \"\"\"formatted output for raster image \n",
    "    \"\"\"\n",
    "    time.sleep(1)\n",
    "    if len(img_data.shape) > 2:\n",
    "        output_img_data = cv2.cvtColor(img_data, cv2.COLOR_BGR2RGB)\n",
    "        output_img_data = output_img_data.transpose(2, 0, 1)\n",
    "        output_full_path = write_path + '/' + prefix + img_name[:-4] + suffix + '.tif'\n",
    "        write_img(output_full_path, im_proj, im_geotrans, output_img_data)\n",
    "    else:\n",
    "        output_full_path = write_path + '/' + prefix + img_name[:-4] + suffix + '.tif'\n",
    "        write_img(output_full_path, im_proj, im_geotrans, img_data)\n",
    "\n",
    "def euclidean_distance(matrix1, matrix2):\n",
    "    \"\"\"calculate euclidean distance\n",
    "    \"\"\"\n",
    "    vector1 = matrix1.flatten().astype(np.float64)\n",
    "    vector2 = matrix2.flatten().astype(np.float64)\n",
    "    \n",
    "    distance = np.sqrt(np.sum((vector1 - vector2) ** 2))\n",
    "    \n",
    "    return distance\n",
    "\n",
    "def average_euclidean_distance(matrix1, matrix2, circle_mask):\n",
    "    \"\"\"calculate average euclidean distance\n",
    "    \"\"\"\n",
    "    bands = np.shape(matrix1)[2]\n",
    "    circle_mask_expand_bands = np.zeros((np.shape(matrix1)[0],np.shape(matrix1)[1],np.shape(matrix1)[2]))\n",
    "    for i in range(bands):\n",
    "        circle_mask_expand_bands[:,:,i] = circle_mask\n",
    "\n",
    "    vector1 = matrix1.flatten().astype(np.float64)\n",
    "    vector2 = matrix2.flatten().astype(np.float64)\n",
    "    vector3 = circle_mask_expand_bands.flatten()\n",
    "    \n",
    "    pixel_distances = vector3 * np.sqrt((vector1 - vector2) ** 2)\n",
    "    \n",
    "    average_distance = np.sum(pixel_distances)/sum(vector3)\n",
    "    \n",
    "    return average_distance\n",
    "\n",
    "def cal_patch_march_num(tmp_labels, circle_mask):\n",
    "    \"\"\"calculate the number of matched pixels\n",
    "    \"\"\"\n",
    "    return np.sum(tmp_labels * circle_mask)\n",
    "\n",
    "def cal_lockmap_match_num(lock_map_patch, circle_mask):\n",
    "    \"\"\"calculate the number of locked pixels\n",
    "    \"\"\"\n",
    "    return np.sum(lock_map_patch * circle_mask)\n",
    "\n",
    "\n",
    "def geo2imagexy(geoX, geoY, im_geotrans):\n",
    "    \"\"\"geo coord to xy coord\n",
    "    \"\"\"\n",
    "    g0 = float(im_geotrans[0])\n",
    "    g1 = float(im_geotrans[1])\n",
    "    g2 = float(im_geotrans[2])\n",
    "    g3 = float(im_geotrans[3])\n",
    "    g4 = float(im_geotrans[4])\n",
    "    g5 = float(im_geotrans[5])\n",
    "\n",
    "    x = (geoX*g5 - g0*g5 - geoX*g2 + g3*g2)/(g1*g5 - g4*g2)\n",
    "    y = (geoY - g3 - geoX*g4)/ g5\n",
    "\n",
    "    return x, y\n",
    "\n",
    "def write_to_log(log_filename, log_message):\n",
    "    \"\"\"write to log file\n",
    "    \"\"\"\n",
    "    with open(log_filename, 'a') as log_file:\n",
    "        log_entry = f'{log_message}'\n",
    "        log_file.write(log_entry)\n",
    "\n",
    "def truncated_linear_stretch(image, truncated_value, max_out = 255, min_out = 0):\n",
    "    \"\"\"image stretching\n",
    "    \"\"\"\n",
    "    def gray_process(gray):\n",
    "        nonzero_pixels = gray[gray != 0]\n",
    "\n",
    "        if len(nonzero_pixels) == 0:\n",
    "            return gray\n",
    "\n",
    "        truncated_down = np.percentile(nonzero_pixels, truncated_value)\n",
    "        truncated_up = np.percentile(nonzero_pixels, 100 - truncated_value)\n",
    "        gray = (gray - truncated_down) / (truncated_up - truncated_down) * (max_out - min_out) + min_out \n",
    "        gray[gray < min_out] = min_out\n",
    "        gray[gray > max_out] = max_out\n",
    "        if(max_out <= 255):\n",
    "            gray = np.uint8(gray)\n",
    "        elif(max_out <= 65535):\n",
    "            gray = np.uint16(gray)\n",
    "        return gray\n",
    "\n",
    "    if(len(image.shape) == 3):\n",
    "        image_stretch = []\n",
    "        for i in range(image.shape[0]):\n",
    "            gray = gray_process(image[i])\n",
    "            image_stretch.append(gray)\n",
    "        image_stretch = np.array(image_stretch)\n",
    "    else:\n",
    "        image_stretch = gray_process(image)\n",
    "    return image_stretch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''parameters'''\n",
    "dataset = 'd2' # testset [d0 or d1 or d2]\n",
    "sr_img_path = r'../data/Test_sets/1-clip_img_' + dataset # image path\n",
    "semantic_segmentation_result_path = r'../data/Test_sets/2-semantic_segmentation_' + dataset # semantic segmentation result path\n",
    "output_shp_path = r'../output/' + dataset # output path\n",
    "template_path = r'../data/Templates' # template path\n",
    "log_file_name = r'../log/log.txt' # log path\n",
    "\n",
    "glt = 0.9 # global lock threshold ***core parameter\n",
    "mlt = 0.8 # minimal lock threshold ***core parameter\n",
    "\n",
    "\n",
    "'''*** NOT recommended to modify the following parameters ***'''\n",
    "foreground_value = 1 # the value of Populus pixels in semantic segmentation result\n",
    "shandow_threshold = 160 # shadow threshold\n",
    "minimum_threshold = 3 # single Populus detection threshold\n",
    "skip_threshold = 0.3 # valid pixel threshold\n",
    "\n",
    "Tempalte = []  # big first\n",
    "#Tempalte.append(np.load(template_path + '/1_19_21.npy'))\n",
    "#Tempalte.append(np.load(template_path + '/1_17_19.npy'))\n",
    "Tempalte.append(np.load(template_path + '/1_15_17.npy')) # [7-17] achieves best performance \n",
    "Tempalte.append(np.load(template_path + '/1_13_15.npy'))\n",
    "Tempalte.append(np.load(template_path + '/1_11_13.npy'))\n",
    "Tempalte.append(np.load(template_path + '/1_9_11.npy'))\n",
    "Tempalte.append(np.load(template_path + '/1_7_9.npy'))\n",
    "Tempalte.append(np.load(template_path + '/1_5_7.npy'))\n",
    "# Tempalte.append(np.load(template_path + '/1_0_5.npy'))\n",
    "\n",
    "T_num = len(Tempalte) # Number of used templates\n",
    "c_size = []\n",
    "for i in range(T_num):\n",
    "    c_size.append(np.shape(Tempalte[i])[0])\n",
    "\n",
    "# mask generation\n",
    "circle_masks = []\n",
    "circle_valid_sums = []\n",
    "\n",
    "for k in range(T_num):\n",
    "    circle_mask = np.zeros((c_size[k], c_size[k]), dtype=np.int32)\n",
    "    circle_radius = c_size[k]//2\n",
    "    circle_center = (c_size[k]//2, c_size[k]//2)\n",
    "    for i in range(c_size[k]):\n",
    "        for j in range(c_size[k]):\n",
    "            distance = np.sqrt((i - circle_center[0])**2 + (j - circle_center[1])**2)\n",
    "            if distance <= circle_radius:\n",
    "                circle_mask[i, j] = 1\n",
    "    circle_masks.append(circle_mask)\n",
    "\n",
    "suit_map_binary_values = []\n",
    "for i in range(T_num): \n",
    "    suit_map_binary_values.append(10)\n",
    "\n",
    "max_distances = [] # the max searching distance of crown center \n",
    "for i in range(T_num): \n",
    "    max_distances.append(c_size[i])\n",
    "\n",
    "lock_thresholds = []\n",
    "for i in range(T_num): \n",
    "    lock_thresholds.append(int(math.pi * (c_size[i]//2) **2) * glt)\n",
    "\n",
    "if T_num>3:\n",
    "    suit_map_binary_values[T_num-1] = 10\n",
    "    circle_masks[T_num-1][:] = 1\n",
    "    lock_thresholds[T_num-1] = c_size[T_num-1] **2 * mlt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''write log'''\n",
    "from datetime import datetime\n",
    "current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
    "write_to_log(log_file_name, current_time)\n",
    "write_to_log(log_file_name, '\\nTmp:'+ str(c_size) +', suit_binaryT:' +str(suit_map_binary_values) +', shadowT:'+str(shandow_threshold)+', minimumT:'+str(minimum_threshold)+', skipT:' + str(skip_threshold)+', maxD:' + str(max_distances) +', lockT:' + str(lock_thresholds) +'\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 43/43 [05:13<00:00,  7.28s/it]\n"
     ]
    }
   ],
   "source": [
    "listpic = fnmatch.filter(os.listdir(sr_img_path), '*.tif')\n",
    "\n",
    "for img_name in tqdm(listpic): # for each test image\n",
    "    img_full_path = sr_img_path + '/' + img_name\n",
    "    sen_seg_full_path = semantic_segmentation_result_path + '/' + img_name\n",
    "\n",
    "    # read image\n",
    "    data_ss, im_proj, im_geotrans = read_img(sen_seg_full_path)[:3]\n",
    "    data_img = read_img(img_full_path)[0]\n",
    "\n",
    "    data_img = cv2.cvtColor(data_img.transpose(1, 2, 0), cv2.COLOR_RGB2BGR) # (c, h, w) -> (h, w, c) & RGB->BGR\n",
    "\n",
    "    # simple preprocessing\n",
    "    height, width, channel = np.shape(data_img)\n",
    "    data_ss = cv2.resize(data_ss, (width, height))\n",
    "\n",
    "    # masking\n",
    "    binary_mask = (data_ss == foreground_value).astype(np.uint8)\n",
    "    masked_img = np.zeros((height, width, channel), dtype=np.uint8)\n",
    "\n",
    "    for i in range(channel):\n",
    "        masked_img[:,:,i] = data_img[:,:,i] * binary_mask\n",
    "\n",
    "    # shadow removing\n",
    "    rt = masked_img[:,:,2].copy()\n",
    "    rt[rt<=shandow_threshold] = 0\n",
    "    rt[rt>shandow_threshold] = 255\n",
    "\n",
    "    # construct connected components in the binary mask\n",
    "    opening = rt.copy()\n",
    "    markers_com = opening.copy()\n",
    "\n",
    "    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(markers_com, connectivity=4)\n",
    "\n",
    "    '''calculate suit map'''\n",
    "    suit_map = np.zeros((T_num, height, width), dtype=np.float64)\n",
    "    suitmap_for_vis = suit_map.copy()\n",
    "    minimum_threshold = minimum_threshold\n",
    "    populus_center = []\n",
    "    skip_threshold = skip_threshold\n",
    "\n",
    "    for label in range(1, num_labels):  # skip background 0\n",
    "        area = stats[label, cv2.CC_STAT_AREA]\n",
    "        x, y, p_width, p_height = stats[label, cv2.CC_STAT_LEFT], stats[label, cv2.CC_STAT_TOP], stats[label, cv2.CC_STAT_WIDTH], stats[label, cv2.CC_STAT_HEIGHT]\n",
    "        \n",
    "        # identify single populus\n",
    "        if p_width<c_size[0] and p_height<c_size[0]:\n",
    "            if p_width > minimum_threshold and p_height > minimum_threshold:\n",
    "                populus_center.append([y+p_height//2, x+p_width//2, -1])\n",
    "                continue\n",
    "            else:\n",
    "                continue\n",
    "\n",
    "        # generate suit map for clustered populus\n",
    "        for k in range(T_num):\n",
    "            for i in range(0, p_height):\n",
    "                for j in range(0, p_width):\n",
    "                    if y+i+c_size[k] < height and x+j+c_size[k] < width:\n",
    "                        tmp_labels = labels[y+i: y+i+c_size[k], x+j: x+j+c_size[k]].copy()\n",
    "                        tmp_labels[tmp_labels!=label] = 0\n",
    "                        tmp_labels[tmp_labels==label] = 1\n",
    "                        if cal_patch_march_num(tmp_labels, circle_masks[k]) >= skip_threshold * sum(circle_masks[k]):\n",
    "                            tmp_img = data_img[y+i: y+i+c_size[k], x+j: x+j+c_size[k], :].copy()\n",
    "                            average_distance = average_euclidean_distance(Tempalte[k], tmp_img, circle_masks[k])\n",
    "                            average_distance = 255-average_distance\n",
    "                            suit_map[k, y+i, x+j] = average_distance\n",
    "                            suitmap_for_vis[k, y+i+c_size[k]//2, x+j+c_size[k]//2] = average_distance\n",
    "                        else:\n",
    "                            continue # skip if the ratio of foreground < threshold\n",
    "                    else:\n",
    "                        continue\n",
    "\n",
    "    '''preprocess for individual trees detection'''\n",
    "    nonzero_values = []\n",
    "    for i in range(T_num):\n",
    "        tmp_suit_map = suit_map[i,:,:]\n",
    "        nonzero_value = tmp_suit_map[tmp_suit_map > 0]\n",
    "        nonzero_values.append(nonzero_value)\n",
    "\n",
    "    threshold_values = []\n",
    "    for i in range(T_num):\n",
    "        if (len(nonzero_values[i]) > 1):\n",
    "            threshold_values.append(np.percentile(nonzero_values[i], suit_map_binary_values[i]))\n",
    "        else:\n",
    "            threshold_values.append(1)\n",
    "\n",
    "    lock_map = opening.copy() # initialize the lock map \n",
    "    lock_map[lock_map==255] = 1    \n",
    "\n",
    "    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)] # searching direction\n",
    "\n",
    "    '''individual trees detection'''\n",
    "    for k in range(T_num):\n",
    "        for i in range(0, height-c_size[k]):\n",
    "            for j in range(0, width-c_size[k]):\n",
    "\n",
    "                visited = np.zeros((height, width), dtype=bool)\n",
    "                max_distance = max_distances[k]\n",
    "                max_value = -999\n",
    "\n",
    "                if suit_map[k, i, j] > threshold_values[k] and not visited[i, j] and lock_map[i,j]==1 and cal_lockmap_match_num(lock_map[i: i+c_size[k], j: j+c_size[k]], circle_masks[k]) >= lock_thresholds[k]:\n",
    "                    # constrained search #1\n",
    "                    queue = deque([(i, j, 0)])\n",
    "                    visited[i][j] = True\n",
    "                    max_value = suit_map[k, i, j]\n",
    "                    max_coords = (i, j)\n",
    "\n",
    "                    while queue:\n",
    "                        h, w, distance = queue.popleft()\n",
    "                        if distance > max_distance:\n",
    "                            continue\n",
    "\n",
    "                        for dh, dw in directions:\n",
    "                            nh, nw = h + dh, w + dw\n",
    "                            if 0 <= nh < height and 0 <= nw < width and suit_map[k, nh, nw] > threshold_values[k] and not visited[nh][nw] and lock_map[nh,nw]==1 and cal_lockmap_match_num(lock_map[nh: nh+c_size[k], nw: nw+c_size[k]], circle_masks[k]) >= lock_thresholds[k]:\n",
    "                                # constrained search #2\n",
    "                                queue.append((nh, nw, distance + 1))\n",
    "                                visited[nh][nw] = True\n",
    "                                if suit_map[k, nh, nw] > max_value:\n",
    "                                    max_value = suit_map[k, nh, nw]\n",
    "                                    max_coords = (nh, nw)\n",
    "                    \n",
    "                    # update lock map\n",
    "                    for m in range(c_size[k]):\n",
    "                        for n in range(c_size[k]):\n",
    "                            if circle_masks[k][m][n] == 1:\n",
    "                                lock_map[max_coords[0]+m, max_coords[1]+n] = 0\n",
    "                    \n",
    "                    # get the crown center if queue is empty\n",
    "                    populus_center.append([max_coords[0]+c_size[k]//2, max_coords[1]+c_size[k]//2, k])   \n",
    "\n",
    "    '''ouput the detection result'''\n",
    "    gdal.SetConfigOption(\"GDAL_FILENAME_IS_UTF8\", \"YES\")\n",
    "    gdal.SetConfigOption(\"SHAPE_ENCODING\", \"GBK\")\n",
    "\n",
    "    ogr.RegisterAll()  # register all dirvers\n",
    "\n",
    "    driver = ogr.GetDriverByName('ESRI Shapefile')\n",
    "\n",
    "    prj = osr.SpatialReference()\n",
    "    prj.ImportFromWkt(im_proj)  # read the projection info from raster image\n",
    "\n",
    "    out_shp_full_path = output_shp_path + '/' +img_name[:-4] + '_populus_infer.shp'\n",
    "    if Path(out_shp_full_path).exists():\n",
    "        driver.DeleteDataSource(out_shp_full_path)\n",
    "    out_ds = driver.CreateDataSource(out_shp_full_path)\n",
    "    out_lyr = out_ds.CreateLayer(\n",
    "        out_shp_full_path, prj, ogr.wkbPoint)\n",
    "    def_out_feature = out_lyr.GetLayerDefn()  # read the feature type\n",
    "\n",
    "    oField = ogr.FieldDefn('type', ogr.OFTInteger)\n",
    "    out_lyr.CreateField(oField)\n",
    "\n",
    "    for i in range(len(populus_center)):\n",
    "        point = ogr.Geometry(ogr.wkbPoint)\n",
    "        x_coord, y_coord = imagexy2geo(im_geotrans, populus_center[i][0], populus_center[i][1])\n",
    "        write_point_to_layer([x_coord, y_coord, populus_center[i][2]], out_lyr, def_out_feature)\n",
    "    out_ds.Destroy()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learn",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
